import time
import mujoco
import mujoco.viewer
import numpy as np

from SDK import SyncTeacherConnection, RealMachineConnection


# def on_joint_angles(degrees):
#     # print('Received joint angles: {}'.format(degrees))
#     global target_positions
#     target_positions = np.deg2rad(degrees)
#     # print('target_positions: {}'.format(target_positions))


if __name__ == "__main__":
    # 加载示教器连接
    teacher = SyncTeacherConnection('config.yml')
    # 加载真实机器人连接
    machine = RealMachineConnection('config.yml')
    
    # Load the model
    model = mujoco.MjModel.from_xml_path('asserts/lift_cube2.xml')
    data = mujoco.MjData(model)

    # 初始化目标位置
    target_positions = np.zeros(6)
    
    # 启动查看器
    with mujoco.viewer.launch_passive(model, data,show_left_ui=False, show_right_ui=False) as viewer:
        # 隐藏左边的控制按钮
        with viewer.lock():
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_ACTUATOR] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CAMERA] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_LIGHT] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONSTRAINT] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_COM] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = 0
            
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_TENDON] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_RANGEFINDER] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONVEXHULL] = 0
        
        while viewer.is_running() :
            step_start = time.time()

            # conn.request_joint_angles(on_joint_angles)
            angles = teacher.request_joint_angles()
            if angles is not None:
                # angles[5] = -angles[5]
                degrees = list(angles)
                degrees[5] = 90 - degrees[5] 
                data.ctrl[:] = np.deg2rad(degrees)
                
                angles[0] = -angles[0]
                angles[1] = -angles[1]
                angles[2] = -angles[2]
                angles[3] = -angles[3]
                angles[4] = -angles[4]
                angles[5] = -angles[5]
                machine.set_joint_angles(angles)
            
            # print(data.ctrl[:])
            # 步进仿真
            mujoco.mj_step(model, data)

            # 同步查看器
            viewer.sync()
            # time.sleep(0.01)